import sys
import tensorflow as tf
from tensorflow.python.platform import gfile
from google.protobuf import text_format

def convert_pb_to_pbtxt(filename):
    with gfile.FastGFile(filename, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        tf.import_graph_def(graph_def, name='')
        tf.train.write_graph(graph_def, './', 'tmp.pbtxt', as_text=True)
    return

def convert_pbtxt_to_pb(filename):
  """Returns a `tf.GraphDef` proto representing the data in the given pbtxt file.
  Args:
    filename: The name of a file containing a GraphDef pbtxt (text-formatted
      `tf.GraphDef` protocol buffer data).
  """
  with tf.gfile.FastGFile(filename, 'r') as f:
    graph_def = tf.GraphDef()
 
    file_content = f.read()
 
    # Merges the human-readable string in `file_content` into `graph_def`.
    text_format.Merge(file_content, graph_def)
    tf.train.write_graph( graph_def , './' , 'tmp.pb' , as_text = False )
  return


if __name__ == '__main__':
    #convert_pbtxt_to_pb(filename)
    #sed -i '/tensor_content/d' protobuf.pbtxt
    #mv protobuf.pbtxt xxx.pbtxt
    #cat xxx.pbtxt | grep op | awk -F '"' '{print$2}' | sort | uniq -c
    filename=sys.argv[1]
    mode=sys.argv[2]
    print(mode) 
    if mode is '0':
        convert_pbtxt_to_pb(filename)
    elif mode == '1':
        convert_pb_to_pbtxt(filename)
    else:
        print("mode input error")
        exit(0)
